package main

import (
	"shell/common"
	"shell/common/packages"
	"shell/global"
	"time"
)

/*
	shell命令请求：非并行，执行后结束
	列出被控端状态请求：非并行、未知
	列出公网端口转发请求：非并行、执行后结束

	端口转发请求：
		本地端口转发：并行、持续
		公网端口转发：非并行、执行后结束
	socks5转发请求：并行、持续
*/

// 处理客户端连接，阻塞
// connCtx是生成quicConn时使用的上下文，用于被动关闭连接
func handleClientEnd(stream *common.QuicStream, metadata *packages.PackageMetadata) {
	action := metadata.PkgAct
	switch action {
	case packages.ActCloseConnReq: // 关闭连接请求
		req := new(packages.CloseConnReq)
		err := stream.ReadPackageAfterDetect(req, metadata)
		if err != nil {
			common.HandleErr(err)
		}
		if req.CloseConn {
			stream.CloseConnSync(Role, true)
			return
		}
		stream.CloseStreamSync(Role, true)
		conn := stream.Conn
		stream = nil
		stream, err := conn.AcceptStream()
		if err != nil {
			common.HandleErr(err)
			return
		}
		metadata, err = common.DetectPackage(stream)
		if err != nil {
			common.HandleErr(err)
			conn.Close(common.CloseByProcess)
			return
		}
		go handleClientEnd(stream, metadata)
	case packages.ActShellCmdReq: // 处理执行shell命令请求
		req := new(packages.ShellCmdReq)
		err := stream.ReadPackageAfterDetect(req, metadata)
		if err != nil {
			common.HandleErr(err)
			stream.CloseConn(common.CloseByProcess)
			return
		}
		err = HandleShellCmdReq(req, stream)
		if err != nil {
			common.HandleErr(err)
			return
		}
		metadata, err = common.DetectPackage(stream)
		if err != nil {
			common.HandleErr(err)
			stream.CloseConn(common.CloseByProcess)
			return
		}
		if metadata.PkgAct != packages.ActCloseConnReq {
			logger.Errorf(0, "error get unexcept package: %s", metadata.PkgAct.String())
			time.Sleep(time.Millisecond * 100)
			stream.CloseConn(common.CloseByProcess)
			return
		}
		go handleClientEnd(stream, metadata)
	case packages.ActPortRedirReq: // 处理端口转发请求
		req := new(packages.PortRedirReq)
		err := stream.ReadPackageAfterDetect(req, metadata)
		if err != nil {
			common.HandleErr(err)
			stream.CloseConn(common.CloseByProcess)
			return
		}
		if req.IsPubNetRedir {
			// 是公网端口转发请求，非持续命令
			err = HandlePubPortRedir(stream, req)
			if err != nil {
				common.HandleErr(err)
				stream.CloseConn(common.CloseByProcess)
				return
			}
			metadata, err = common.DetectPackage(stream)
			if err != nil {
				common.HandleErr(err)
				stream.CloseConn(common.CloseByProcess)
				return
			}
			if metadata.PkgAct != packages.ActCloseConnReq {
				logger.Errorf(0, "error get unexcept package: %s", metadata.PkgAct.String())
				time.Sleep(time.Millisecond * 200)
				stream.CloseConn(common.CloseByProcess)
				return
			}
			go handleClientEnd(stream, metadata)
		} else {
			// 是持续性转发请求，和socks5转发类似
			go HandlePrivatePortRedirConn(stream.Conn)
			HandlePrivatePortRedirStream(stream, req, nil)
		}
	case packages.ActSocks5RedirReq: // 处理socks5转发请求
		go HandleSocks5Connection(stream.Conn)
		HandleSingleSocks5Stream(stream, metadata)

	case packages.ActControlledReq: // 处理列出被控端请求
		req := new(packages.ControlledReq)
		err := stream.ReadPackageAfterDetect(req, metadata)
		if err != nil {
			common.HandleErr(err)
			stream.CloseConn(common.CloseByProcess)
			return
		}
		resp := new(packages.ControlledResp)
		resp.CmdUUID = req.CmdUUID
		resp.SetRole(Role)
		resp.Stats = ListControlledEnd()
		err = stream.WritePackage(resp)
		if err != nil {
			common.HandleErr(err)
			stream.CloseConn(common.CloseByProcess)
			return
		}
		metadata, err = common.DetectPackage(stream)
		if err != nil {
			common.HandleErr(err)
			stream.CloseConn(common.CloseByProcess)
			return
		}
		go handleClientEnd(stream, metadata)
	case packages.ActListGWPortReq: // 处理列出公网端口转发请求
		req := new(packages.ListGWPortReq)
		err := stream.ReadPackageAfterDetect(req, metadata)
		if err != nil {
			common.HandleErr(err)
			stream.CloseConn(common.CloseByProcess)
			return
		}
		err = HandleListPort(req, stream)
		if err != nil {
			common.HandleErr(err)
			return
		}
		metadata, err = common.DetectPackage(stream)
		if err != nil {
			common.HandleErr(err)
			stream.CloseConn(common.CloseByProcess)
			return
		}
		if metadata.PkgAct != packages.ActCloseConnReq {
			logger.Errorf(0, "error get unexcept package: %s", metadata.PkgAct.String())
			time.Sleep(time.Millisecond * 200)
			stream.CloseConn(common.CloseByProcess)
			return
		}
		go handleClientEnd(stream, metadata)
	default:
		// 非法请求，关闭连接并退出
		logger.Errorf(0, "error unsupport package action: %d", action)
		stream.CloseConn(common.CloseByProcess)
		return
	}

}

// 处理shell命令请求，如果发生错误，会关闭流和quic连接
func HandleShellCmdReq(req *packages.ShellCmdReq, stream *common.QuicStream) error {
	val, have := ControlledEndMap.Load(req.TargetUUID)
	resp := new(packages.ShellCmdResp)
	resp.FromReq(req, Role)
	if have {
		// 在全局map中有指定的被控端
		cei, ok := val.(*ControlledEndInfo)
		if ok {
			if cei.IsProxy() {
				// 目标节点是公网转发端，直接执行命令即可
				return HandleShellCmdReqToProxy(req, stream)
			}
			// 获取目标节点的流
			cStream, err := cei.OpenStream()
			if err != nil {
				// 打印错误
				common.HandleErr(err)
				resp.CmdExitCode = -1
				resp.CmdOutput = []byte(err.Error())
				errW := stream.WritePackage(resp)
				if errW != nil {
					stream.CloseConn(common.CloseByProcess)
				}
				return errW
			}
			// 以转发端身份转发shell请求
			req.SetRole(Role)
			err = cStream.WritePackage(req)
			if err != nil {
				// 打印错误
				common.HandleErr(err)
				resp.CmdExitCode = -1
				resp.CmdOutput = []byte(err.Error())
				errW := stream.WritePackage(resp)
				if errW != nil {
					stream.CloseConn(common.CloseByProcess)
				}
				return errW
			}
			err = cStream.ReadPackage(resp)
			resp.SetRole(Role)
			if err != nil {
				resp.CmdExitCode = -1
				resp.CmdOutput = []byte(err.Error())
				errW := stream.WritePackage(resp)
				if errW != nil {
					stream.CloseConn(common.CloseByProcess)
				}
				return errW
			}
			err = stream.WritePackage(resp)
			if err != nil {
				stream.CloseConn(common.CloseByProcess)
			}
			return err
		}
	}
	// 没有在全局map中查询到指定的被控端
	resp.CmdExitCode = -1
	resp.CmdOutput = []byte("target controlled end is down")
	err := stream.WritePackage(resp)
	if err != nil {
		stream.CloseConn(common.CloseByProcess)
	}
	return err
}

// 处理列出端口请求，如果发生错误，就关闭stream流和连接
func HandleListPort(req *packages.ListGWPortReq, stream *common.QuicStream) error {
	resp := new(packages.ListGWPortResp)
	resp.SetRole(packages.RoleProxy)
	if globalCfg.TcpPorts != nil && globalCfg.TcpPorts.Enable {
		resp.Enabled = true
		resp.StartPort = globalCfg.TcpPorts.From
		resp.EndPort = globalCfg.TcpPorts.To
	} else {
		resp.Enabled = false
		resp.StartPort = 0
		resp.EndPort = 0
	}
	resp.CmdUUID = req.CmdUUID
	resp.Stats = ListPortRedir()
	return stream.WritePackage(resp)
}

// 处理socks5请求，是阻塞方法
func HandleSocks5Connection(quicConn *common.QuicConn) {
	for {
		stream, err := quicConn.AcceptStream()
		if err != nil {
			common.HandleErr(err)
			break
		}
		md, err := stream.DetectPackage()
		if err != nil {
			common.HandleErr(err)
			continue
		}
		switch md.PkgAct {
		case packages.ActSocks5RedirReq:
			go HandleSingleSocks5Stream(stream, md)
		case packages.ActKeepaliveReq:
			go HandleKeepaliveReq(stream, md)
		default:
			logger.Errorf(0, "error get package type: %s", md.PkgAct)
			stream.Close()
		}
	}
}

func HandleKeepaliveReq(stream *common.QuicStream, metadata *packages.PackageMetadata) {
	req := new(packages.KeepaliveReq)
	err := stream.ReadPackageAfterDetect(req, metadata)
	if err != nil {
		common.HandleErr(err)
		return
	}
	resp := new(packages.KeepaliveResp)
	resp.SetRole(Role)
	resp.CmdUUID = req.CmdUUID
	err = stream.WritePackage(resp)
	if err != nil {
		common.HandleErr(err)
		return
	}
	stream.WaitForCloseSync(Role, false)
}

// 阻塞，处理单个socks5转发请求，如果metadata为nil，表示由HandleSocks5Connection方法调用
func HandleSingleSocks5Stream(stream *common.QuicStream, metadata *packages.PackageMetadata) {
	req := new(packages.Socks5RedirReq)
	if metadata != nil {
		// 已经探明是socks5转发请求
		err := stream.ReadPackageAfterDetect(req, metadata)
		if err != nil {
			common.HandleErr(err)
			return
		}
	} else {
		err := stream.ReadPackage(req)
		if err != nil {
			common.HandleErr(err)
			return
		}
	}
	respFunc := func(errInfo string) {
		resp := req.ToResp(Role)
		resp.IsOk = false
		resp.ErrorInfo = []byte(errInfo)
		err := stream.WritePackage(resp)
		if err != nil {
			common.HandleErr(err)
			return
		}
		stream.WaitForClose(Role, false)
	}
	val, ok := ControlledEndMap.Load(req.TargetUUID)
	if !ok {
		logger.Errorf(0, "target controlled end is down")
		respFunc("target controlled end is down")
		return
	}
	cei := val.(*ControlledEndInfo)
	if cei.IsProxy() {
		// socks5转发的目的节点是公网转发端
		HandleSingleSocks5StreamToProxy(stream, req)
		return
	}
	cStream, err := cei.OpenStream()
	if err != nil {
		common.HandleErr(err)
		respFunc("cat't open a stream to target controlled end: " + err.Error())
		return
	}
	req.SetRole(Role)
	err = cStream.WritePackage(req)
	if err != nil {
		common.HandleErr(err)
		respFunc("error send socks5 redir req to target controlled end: " + err.Error())
		return
	}
	resp := new(packages.Socks5RedirResp)
	err = cStream.ReadPackage(resp)
	if err != nil {
		common.HandleErr(err)
		respFunc("error read socks5 redir resp from target controlled end: " + err.Error())
		return
	}
	if !resp.IsOk {
		logger.Infof(0, "get error from target controlled end: %s", string(resp.ErrorInfo))
		respFunc("get error from target controlled end: " + string(resp.ErrorInfo))
		cStream.WaitForClose(Role, true)
		return
	}
	resp.CmdUUID = req.CmdUUID
	resp.SetRole(Role)
	resp.IsOk = true
	// 向命令行端回复socks5转发响应
	err = stream.WritePackage(resp)
	if err != nil {
		common.HandleErr(err)
		cStream.Close()
		return
	}
	common.NewTimeoutRWCPairWarp(stream, cStream, global.DefaultIdleTime, timeoutWatcher).RedirSyncNoKey()
}

// 非阻塞，处理公网端口转发请求，如果发生错误就关闭流和连接
func HandlePubPortRedir(stream *common.QuicStream, req *packages.PortRedirReq) error {
	if req == nil {
		stream.CloseConn(common.CloseByProcess)
		return global.ErrNil
	}
	pi, err := GetOrCreatePortInfo(req.RedirSrcPort)
	resp := req.ToResp(packages.RoleProxy)
	if err != nil {
		resp.IsOk = false
		resp.ErrorInfo = []byte(err.Error())
		errW := stream.WritePackage(resp)
		if errW != nil {
			common.HandleErr(errW)
			stream.Close()
			return errW
		}
		return nil
	}
	if req.Start {
		if pi.Stat.Load() != int32(packages.PortHardStop) {
			// 状态错误
			resp.IsOk = false
			resp.ErrorInfo = []byte("public network port status is not hardstop, can't start now")
			errW := stream.WritePackage(resp)
			if errW != nil {
				common.HandleErr(errW)
				stream.Close()
				return errW
			}
			return nil
		}
		err = pi.Start(req)
		if err != nil {
			resp.IsOk = false
			resp.ErrorInfo = []byte("public network port start fail: " + err.Error())
			errW := stream.WritePackage(resp)
			if errW != nil {
				common.HandleErr(errW)
				stream.Close()
				return errW
			}
			return nil
		}
	}
	if req.HardStop {
		pi.HardStop()
	}
	if req.SoftStop {
		pi.SoftStop()
	}
	resp.IsOk = true
	err = stream.WritePackage(resp)
	if err != nil {
		common.HandleErr(err)
		stream.Close()
		return err
	}
	return nil
}

// 处理单个端口转发流，如果req不为nil，表示已经接收了请求
func HandlePrivatePortRedirStream(stream *common.QuicStream, req *packages.PortRedirReq, md *packages.PackageMetadata) {
	if md != nil && req == nil {
		req = new(packages.PortRedirReq)
		err := stream.ReadPackageAfterDetect(req, md)
		if err != nil {
			common.HandleErr(err)
			return
		}
	}
	if req == nil {
		req = new(packages.PortRedirReq)
		err := stream.ReadPackage(req)
		if err != nil {
			common.HandleErr(err)
			return
		}
	}
	respFunc := func(errInfo string) {
		resp := req.ToResp(packages.RoleProxy)
		resp.IsOk = false
		resp.ErrorInfo = []byte(errInfo)
		err := stream.WritePackage(resp)
		if err != nil {
			common.HandleErr(err)
			return
		}
		stream.WaitForClose(Role, false)
	}
	redirRule, err := req.CheckAndClone(false, Role)
	if err != nil {
		common.HandleErr(err)
		respFunc(err.Error())
		return
	}
	redirRule.SetRole(packages.RoleProxy)
	val, ok := ControlledEndMap.Load(redirRule.TargetUUID)
	if !ok {
		respFunc("target controlled end is down")
		return
	}
	cei := val.(*ControlledEndInfo)
	if cei.IsProxy() {
		// 端口转发的目标节点是转发端
		HandlePrivatePortRedirStreamToProxy(stream, req)
		return
	}
	cStream, err := cei.OpenStream()
	if err != nil {
		respFunc("open stream fail: " + err.Error())
		return
	}

	err = cStream.WritePackage(redirRule)
	if err != nil {
		logger.Errorf(0, "error write port redir req to controlled end: %v", err)
		respFunc("error write port redir req to controlled end: " + err.Error())
		return
	}
	resp := new(packages.PortRedirResp)
	err = cStream.ReadPackage(resp)
	if err != nil {
		logger.Errorf(0, "error read port redir resp to controlled end: %v", err)
		respFunc("error read port redir resp to controlled end: " + err.Error())
		return
	}
	if !resp.IsOk {
		respFunc("error from controlled end: " + string(resp.ErrorInfo))
		return
	}
	resp.SetRole(Role)
	err = stream.WritePackage(resp)
	if err != nil {
		common.HandleErr(err)
		cStream.Close()
		return
	}
	common.NewTimeoutRWCPairWarp(stream, cStream, global.DefaultIdleTime, timeoutWatcher).RedirSyncNoKey()
}

// 处理端口转发quic连接，结束时需要调用ctx中的cancel
func HandlePrivatePortRedirConn(quicConn *common.QuicConn) {
	for {
		stream, err := quicConn.AcceptStream()
		if err != nil {
			common.HandleErr(err)
			break
		}
		md, err := stream.DetectPackage()
		if err != nil {
			common.HandleErr(err)
			continue
		}
		switch md.PkgAct {
		case packages.ActKeepaliveReq:
			// debug
			logger.Debugf(0, "get keep req")
			go HandleKeepaliveReq(stream, md)
		case packages.ActPortRedirReq:
			go HandlePrivatePortRedirStream(stream, nil, md)
		default:
			logger.Errorf(0, "error get package type: %s", md.PkgAct)
			stream.Close()
		}
	}
}
